/*
* KiPE.cpp
*
*  Created on: 2013-10-16
*      Author: WinDDK
*/

#include "KiPE.h"

KiPE::KiPE()
{
    KiDebug("Success");
    RtlZeroMemory(mModuleName, 0, MODULENAMELENGTH + 1);
    RtlZeroMemory(mSystemPath, 0, SYSTEMPATHLENGTH + 1);
    RtlZeroMemory(mModulePath, 0, MODULEPATHLENGTH + 1);

    mNewImageBase = NULL;
    mNewImageNtHeader = NULL;
}

KiPE::~KiPE()
{
    if (mNewImageBase != NULL)
    {
        ExFreePool((PVOID)mNewImageBase);
    }
    KiDebug("Success");
}

BOOLEAN KiPE::LoadModule()
{
    if (!GetModulePath())
    {
        KiDebug("GetModulePath Failed");
        return FALSE;
    }

    if (!MapFileToMem())
    {
        KiDebug("MapFileToMem Failed");
        return FALSE;
    }

    if (!GetImageBase(mModuleName, mOldImageBase))
    {
        KiDebug("GetOldImageBase Failed");
        return FALSE;
    }
    return TRUE;
}

BOOLEAN KiPE::GetImageDosHeader(IN ULONG imageBase, OUT PIMAGE_DOS_HEADER &imageDosHeader)
{
    imageDosHeader = (PIMAGE_DOS_HEADER) imageBase;
    if (imageBase == (ULONG)NULL || imageDosHeader->e_magic != IMAGE_DOS_SIGNATURE)
    {
        KiDebug("GetImageDosHeader Failed");
        imageDosHeader = 0;
        return FALSE;
    }
    return TRUE ;
}

BOOLEAN KiPE::GetImageNTHeader(IN ULONG imageBase, OUT PIMAGE_NT_HEADERS & imageNtHeader)
{
    PIMAGE_DOS_HEADER imageDosHeader;

    if (!GetImageDosHeader(imageBase, imageDosHeader))
    {
        KiDebug("GetImageNTHeader Failed");
        return FALSE;
    }

    imageNtHeader = (PIMAGE_NT_HEADERS) (imageBase + imageDosHeader->e_lfanew);

    if (imageNtHeader->Signature != IMAGE_NT_SIGNATURE)
    {
        KiDebug("GetImageNTHeader Failed");
        imageNtHeader = 0;
        return FALSE;
    }
    return TRUE;
}

BOOLEAN KiPE::MapFileToMem()
{
    ANSI_STRING modulePathAString;
    UNICODE_STRING modulePathUString;

    RtlInitAnsiString(&modulePathAString, mModulePath);
    RtlAnsiStringToUnicodeString(&modulePathUString, &modulePathAString, TRUE);

    HANDLE file;
    OBJECT_ATTRIBUTES objAttr;
    IO_STATUS_BLOCK ioStatusBlock;
    InitializeObjectAttributes(&objAttr, &modulePathUString, OBJ_CASE_INSENSITIVE| OBJ_KERNEL_HANDLE, NULL, NULL);

    NTSTATUS status = ZwCreateFile(&file, FILE_EXECUTE | SYNCHRONIZE, &objAttr, &ioStatusBlock, 0, FILE_ATTRIBUTE_NORMAL, FILE_SHARE_READ,
        FILE_OPEN, FILE_NON_DIRECTORY_FILE | FILE_RANDOM_ACCESS | FILE_SYNCHRONOUS_IO_NONALERT, NULL, 0);
    if (!NT_SUCCESS(status))
    {
        KiDebug("ZwCreateFile Failed");
        return FALSE;
    }

    IMAGE_DOS_HEADER imageDosHeader;
    ULONG imageDosHeaderSize = sizeof(IMAGE_DOS_HEADER);
    LARGE_INTEGER fileOffset = { 0 };
    status = ZwReadFile(file, NULL, NULL, NULL, &ioStatusBlock, &imageDosHeader, imageDosHeaderSize, &fileOffset, 0);
    if (!NT_SUCCESS(status) || imageDosHeader.e_magic != IMAGE_DOS_SIGNATURE)
    {
        KiDebug("Read ImageDosHeader Failed");
        ZwClose(file);
        return FALSE;
    }

    IMAGE_FILE_HEADER imageFileHeader;
    ULONG imageFileHeaderSize = sizeof(IMAGE_FILE_HEADER);
    fileOffset.QuadPart = imageDosHeader.e_lfanew + sizeof(ULONG);
    status = ZwReadFile(file, NULL, NULL, NULL, &ioStatusBlock, &imageFileHeader, imageFileHeaderSize, &fileOffset, 0);
    if (!NT_SUCCESS(status) || imageDosHeader.e_magic != IMAGE_DOS_SIGNATURE)
    {
        KiDebug("Read ImageFileHeader Failed");
        ZwClose(file);
        return FALSE;
    }

    IMAGE_NT_HEADERS imageNtHeader;
    ULONG imageNtHeaderSize = sizeof(ULONG) + imageFileHeaderSize + imageFileHeader.SizeOfOptionalHeader;
    fileOffset.QuadPart = imageDosHeader.e_lfanew;
    status = ZwReadFile(file, NULL, NULL, NULL, &ioStatusBlock, &imageNtHeader, imageNtHeaderSize, &fileOffset, 0);
    if (!NT_SUCCESS(status) || imageNtHeader.Signature != IMAGE_NT_SIGNATURE)
    {
        KiDebug("Read ImageNtHeaders Failed");
        ZwClose(file);
        return FALSE;
    }

    ULONG imageSectionHeaderSize = sizeof(IMAGE_SECTION_HEADER) * imageNtHeader.FileHeader.NumberOfSections;
    IMAGE_SECTION_HEADER *imageSectionHeader = (IMAGE_SECTION_HEADER *) ExAllocatePoolWithTag(NonPagedPool, imageSectionHeaderSize,
            KIPETAG);
    fileOffset.QuadPart = imageDosHeader.e_lfanew + imageNtHeaderSize;
    status = ZwReadFile(file, NULL, NULL, NULL, &ioStatusBlock, imageSectionHeader, imageSectionHeaderSize, &fileOffset, 0);
    if (!NT_SUCCESS(status))
    {
        KiDebug("Read ImageSectionHeader Failed");
        ExFreePool(imageSectionHeader);
        ZwClose(file);
        return FALSE;
    }

    PVOID imageBase = ExAllocatePoolWithTag(NonPagedPool, imageNtHeader.OptionalHeader.SizeOfImage, KIPETAG);
    if (imageBase == 0)
    {
        KiDebug("ExAllocatePoolWithTag Failed");
        ExFreePool(imageSectionHeader);
        ZwClose(file);
        return FALSE;
    }
    RtlZeroMemory(imageBase, imageNtHeader.OptionalHeader.SizeOfImage);

    ULONG headersSize = imageNtHeader.OptionalHeader.SizeOfHeaders;
    fileOffset.QuadPart = 0;
    status = ZwReadFile(file, NULL, NULL, NULL, &ioStatusBlock, imageBase, headersSize, &fileOffset, 0);
    if (!NT_SUCCESS(status))
    {
        KiDebug("Read PE Headers Failed");
        ExFreePool(imageSectionHeader);
        ExFreePool(imageBase);
        ZwClose(file);
        return FALSE;
    }

    for (ULONG i = 0; i < imageNtHeader.FileHeader.NumberOfSections; i++)
    {
        ULONG fileSectionSize = imageSectionHeader[i].SizeOfRawData;
        CHAR sectionName[IMAGE_SIZEOF_SHORT_NAME + 1] = { 0 };
        strncpy(sectionName, (PCHAR)imageSectionHeader[i].Name, IMAGE_SIZEOF_SHORT_NAME);
        KiDebug("%-8s, VS:0x%05x, RS:0x%05x", sectionName, imageSectionHeader[i].Misc.VirtualSize, fileSectionSize);
        fileOffset.QuadPart = imageSectionHeader[i].PointerToRawData;
        status = ZwReadFile(file, NULL, NULL, NULL, &ioStatusBlock, (PVOID) ((ULONG) imageBase + imageSectionHeader[i].VirtualAddress),
            fileSectionSize, &fileOffset, 0);
        if (!NT_SUCCESS(status))
        {
            KiDebug("SectionData Read Failed");
            ExFreePool(imageSectionHeader);
            ExFreePool(imageBase);
            ZwClose(file);
            return FALSE;
        }
    }

    ExFreePool(imageSectionHeader);

    this->mNewImageBase = (ULONG)imageBase;
    this->mNewImageNtHeader = (PIMAGE_NT_HEADERS) ((ULONG) imageBase + imageDosHeader.e_lfanew);

    KiDebug("%s new image base:0x%x", mModuleName, imageBase);
    ZwClose(file);

    return TRUE;
}

BOOLEAN KiPE::GetImageDataDirectory(IN ULONG imageBase, IN ULONG TYPE, OUT PULONG pVA, OUT PULONG pSize)
{
    PIMAGE_NT_HEADERS imageNtHeader;

    if(!GetImageNTHeader(imageBase, imageNtHeader))
    {
        KiDebug("GetImageNTHeader Failed");
    }

    IMAGE_DATA_DIRECTORY ImageDataDirectory = imageNtHeader->OptionalHeader.DataDirectory[TYPE];
    *pVA = (ULONG)imageBase + ImageDataDirectory.VirtualAddress;
    *pSize = ImageDataDirectory.Size;

    if (!MmIsAddressValid((PVOID) *pVA))
    {
        KiDebug("GetImageDataDirectory Failed");
        *pVA = 0;
        *pSize = 0;
        return FALSE;
    }

    return TRUE;
}

BOOLEAN KiPE::GetSystemPath(OUT PCHAR systemPath)
{
    UNICODE_STRING dirName;
    OBJECT_ATTRIBUTES dirObject;
    HANDLE dirHandle;
    RtlInitUnicodeString(&dirName, L"\\KnownDlls");
    InitializeObjectAttributes(&dirObject, &dirName, OBJ_CASE_INSENSITIVE|OBJ_KERNEL_HANDLE, NULL, NULL);
    NTSTATUS status = ZwOpenDirectoryObject(&dirHandle, DIRECTORY_QUERY, &dirObject);
    if (!NT_SUCCESS(status))
    {
        KiDebug("ZwOpenDirectoryObject Failed");
        ZwClose(dirHandle);
        return false;
    }

    UNICODE_STRING symbolName;
    OBJECT_ATTRIBUTES symbolObject;
    HANDLE symbolHandle;
    RtlInitUnicodeString(&symbolName, L"KnownDllPath");
    InitializeObjectAttributes(&symbolObject, &symbolName, OBJ_CASE_INSENSITIVE|OBJ_KERNEL_HANDLE, dirHandle, NULL);
    status = ZwOpenSymbolicLinkObject(&symbolHandle, GENERIC_READ, &symbolObject);
    if (!NT_SUCCESS(status))
    {
        KiDebug("ZwOpenSymbolicLinkObject Failed");
        ZwClose(dirHandle);
        ZwClose(symbolHandle);
        return false;
    }

    WCHAR WCharBuffer[256];
    UNICODE_STRING symbolicUString;
    symbolicUString.Buffer = WCharBuffer;
    symbolicUString.MaximumLength = 256 * sizeof(WCHAR);
    symbolicUString.Length = 0;
    status = ZwQuerySymbolicLinkObject(symbolHandle, &symbolicUString, NULL);
    if (!NT_SUCCESS(status))
    {
        KiDebug("ZwQuerySymbolicLinkObject Failed");
        ZwClose(dirHandle);
        ZwClose(symbolHandle);
        return false;
    }

    ANSI_STRING SymbolicAString;
    status = RtlUnicodeStringToAnsiString(&SymbolicAString, &symbolicUString, TRUE);
    if (!NT_SUCCESS(status))
    {
        KiDebug("RtlUnicodeStringToAnsiString Failed");
        ZwClose(dirHandle);
        ZwClose(symbolHandle);
        return false;
    }

    strncpy(systemPath, SymbolicAString.Buffer, SYSTEMPATHLENGTH);
    KiDebug("%s", systemPath);
    return true;
}

BOOLEAN KiPE::GetImageBase(IN PCHAR moduleName, OUT ULONG &imageBase)
{
    ULONG size;
    NTSTATUS status = ZwQuerySystemInformation(SystemModuleInformation, NULL, 0, &size);
    if (status != STATUS_INFO_LENGTH_MISMATCH )
    {
        KiDebug("ZwQuerySystemInformation Failed");
        return FALSE;
    }
    PSYSTEM_MODULES modules = (PSYSTEM_MODULES) ExAllocatePoolWithTag(NonPagedPool, size, KIPETAG);
    if (modules)
    {
        status = ZwQuerySystemInformation(SystemModuleInformation, modules, size, &size);
        if (!NT_SUCCESS(status))
        {
            KiDebug("ZwQuerySystemInformation Failed");
            ExFreePool((PVOID) modules);
            return FALSE;
        }
    } else
    {
        KiDebug("ExAllocatePoolWithTag Failed");
        return FALSE;
    }
    for (ULONG i = 0; i < modules->NumberOfModules; i++)
    {
        PCHAR name = modules->Modules[i].ImageName + modules->Modules[i].ModuleNameOffset;
        if (!strcmp(_strlwr(name), _strlwr(moduleName)))
        {
            imageBase = (ULONG) modules->Modules[i].ImageBase;
            KiDebug("%s image base:0x%x", name, imageBase);
            ExFreePool(modules);
            return TRUE;
        }
    }
    imageBase = 0;
    KiDebug("Find %s image base Failed", moduleName);
    ExFreePool(modules);
    return FALSE;
}

BOOLEAN KiPE::FixRelocTable()
{
    PIMAGE_BASE_RELOCATION imageBaseRelocation;
    ULONG relocationSize;

    if (!GetImageDataDirectory(mNewImageBase, IMAGE_DIRECTORY_ENTRY_BASERELOC, (PULONG) &imageBaseRelocation,
        &relocationSize))
    {
        KiDebug("GetModuleDataDirectory Failed");
        return FALSE;
    }

    while (relocationSize > 0 && imageBaseRelocation->SizeOfBlock > 0)
    {
        PUSHORT relocationArray = (USHORT*) ((ULONG) imageBaseRelocation + sizeof(IMAGE_BASE_RELOCATION));
        ULONG Num = (imageBaseRelocation->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(USHORT);
        for (ULONG i = 0; i < Num; i++)
        {
            if (relocationArray[i] && (relocationArray[i] >> 12) == IMAGE_REL_BASED_HIGHLOW)
            {
                ULONG relocationAddress = mNewImageBase + imageBaseRelocation->VirtualAddress + (relocationArray[i] & 0xFFF);
                if (*(ULONG*) relocationAddress != 0xcccccccc)
                {
                    *(ULONG*) relocationAddress +=  mOldImageBase - mNewImageNtHeader->OptionalHeader.ImageBase;
                    //*(ULONG*) relocationAddress +=  mNewImageBase - mNewImageNtHeader->OptionalHeader.ImageBase;
                }
            }
        }
        relocationSize -= imageBaseRelocation->SizeOfBlock;
        imageBaseRelocation = (PIMAGE_BASE_RELOCATION) ((ULONG) imageBaseRelocation
            + imageBaseRelocation->SizeOfBlock);
    }
    return TRUE ;
}

BOOLEAN KiPE::FixImportTable()
{
    PIMAGE_IMPORT_DESCRIPTOR imageImportDescriptor;
    ULONG importSize;
    PIMAGE_IMPORT_BY_NAME imortByName;
    ULONG functionAddress;

    if (!GetImageDataDirectory(mNewImageBase, IMAGE_DIRECTORY_ENTRY_IMPORT, (PULONG) &imageImportDescriptor,
        &importSize))
    {
        KiDebug("GetModuleDataDirectory Failed");
        return FALSE;
    }

    while (imageImportDescriptor->OriginalFirstThunk && imageImportDescriptor->Name)
    {
        char moduleName[260];
        strcpy(moduleName, (char*) (mNewImageBase + imageImportDescriptor->Name));
        ULONG imageBase;
        if (!GetImageBase(moduleName, imageBase))
        {
            KiDebug("FixKernelImportTable Failed");
            return FALSE;
        }

        if (imageBase == (ULONG) NULL )
        {
            KiDebug("FixModuleImportTable Alarm!");
            PIMAGE_THUNK_DATA FirstThunk = (PIMAGE_THUNK_DATA) (mNewImageBase + imageImportDescriptor->FirstThunk);
            PIMAGE_THUNK_DATA OldFirstThunk = (PIMAGE_THUNK_DATA) ((ULONG) FirstThunk + mOldImageBase - mNewImageBase);
            while (OldFirstThunk->u1.Function)
            {
                FirstThunk->u1.Function = OldFirstThunk->u1.Function;
                OldFirstThunk++;
                FirstThunk++;
            }
            imageImportDescriptor++;
            continue;
        }
        PIMAGE_THUNK_DATA ImageThunkData = (PIMAGE_THUNK_DATA) (mNewImageBase + imageImportDescriptor->OriginalFirstThunk);
        PIMAGE_THUNK_DATA FirstThunk = (PIMAGE_THUNK_DATA) (mNewImageBase + imageImportDescriptor->FirstThunk);
        KiDebug("FirstThunk address is:0x%x", FirstThunk);
        while (ImageThunkData->u1.Ordinal)
        {
            //序号导入
            if (IMAGE_SNAP_BY_ORDINAL32(ImageThunkData->u1.Ordinal))
            {
                if (!GetExportedRoutineByOrdinal(imageBase, ImageThunkData->u1.Ordinal & ~IMAGE_ORDINAL_FLAG32,
                    &functionAddress))
                {
                    KiDebug("GetExportedRoutineByOrdinal Failed");
                    return FALSE;
                }

                if (functionAddress == 0)
                {
                    KiDebug("can't find funcion Index %d \n",ImageThunkData->u1.Ordinal & ~IMAGE_ORDINAL_FLAG32);
                    break;
                }
                FirstThunk->u1.Function = functionAddress;
            }
            //函数名导入
            else
            {
                imortByName = (PIMAGE_IMPORT_BY_NAME) (mNewImageBase + ImageThunkData->u1.AddressOfData);
                if (!GetExportedRoutineByName(imageBase, (PCHAR)imortByName->Name, &functionAddress))
                {
                    KiDebug("GetExportedRoutineByName Failed");
                    return FALSE;
                }

                if (functionAddress == 0)
                {
                    KiDebug("can't Funcion Name:%s\n",imortByName->Name);
                    break;
                }
                FirstThunk->u1.Function = functionAddress;
            }
            FirstThunk++;
            ImageThunkData++;
        }
        imageImportDescriptor++;
    }
    return TRUE;
}

BOOLEAN KiPE::GetExportedRoutineByOrdinal(ULONG imageBase, ULONG Ordinal, PULONG pFunctionAddress)
{
    PIMAGE_EXPORT_DIRECTORY pImageExportDirectory;
    ULONG ExportSize;
    if (!GetImageDataDirectory(imageBase, IMAGE_DIRECTORY_ENTRY_EXPORT, (PULONG) &pImageExportDirectory,
        (PULONG) &ExportSize))
    {
        KiDebug("GetModuleDataDirectory Failed");
        return FALSE;
    }

    PULONG FunctionsArray = (PULONG) (imageBase + pImageExportDirectory->AddressOfFunctions);
    PULONG NamesArray = (PULONG) (imageBase + pImageExportDirectory->AddressOfNames);
    PUSHORT NameOrdinalsArray = (PUSHORT) (imageBase + pImageExportDirectory->AddressOfNameOrdinals);

    if (Ordinal < pImageExportDirectory->Base
        || Ordinal > pImageExportDirectory->NumberOfFunctions - 1 + pImageExportDirectory->Base )
    {
        KiDebug("GetExportedRoutineByOrdinal Failed");
        *pFunctionAddress = (ULONG) NULL;
        return FALSE;
    }

    ULONG i;
    for (i = 0; i < pImageExportDirectory->NumberOfNames; i++)
    {
        if(Ordinal ==  pImageExportDirectory->Base + NameOrdinalsArray[i])
        {
            break;
        }
    }

    if (i > pImageExportDirectory->NumberOfNames - 1)
    {
        KiDebug("GetExportedRoutineByOrdinal Failed");
        *pFunctionAddress = 0;
        return FALSE;
    }


    char * RoutineName = (char *)((ULONG)imageBase + NamesArray[i]);
    ULONG functionAddress = imageBase + FunctionsArray[NameOrdinalsArray[i]];

    KiDebug("Ordinal:%2d Hint:%2d Function:%s Entry Point:0x%08x", Ordinal, i, RoutineName, FunctionsArray[NameOrdinalsArray[i]]);

    if (!MmIsAddressValid((PVOID)functionAddress))
    {
        KiDebug("Address:0x%x is not Valid", functionAddress);
    }

    *pFunctionAddress = functionAddress;
    return TRUE ;
}

BOOLEAN KiPE::GetExportedRoutineByName(ULONG imageBase, char *RoutineName, PULONG pFunctionAddress)
{
    PIMAGE_EXPORT_DIRECTORY pImageExportDirectory;
    ULONG ExportSize;
    if (!GetImageDataDirectory(imageBase, IMAGE_DIRECTORY_ENTRY_EXPORT, (PULONG) &pImageExportDirectory,
        (PULONG) &ExportSize))
    {
        KiDebug("GetModuleDataDirectory Failed");
        return FALSE;
    }

    PULONG FunctionsArray = (PULONG) (imageBase + pImageExportDirectory->AddressOfFunctions);
    PULONG NamesArray = (PULONG) (imageBase + pImageExportDirectory->AddressOfNames);
    PUSHORT NameOrdinalsArray = (PUSHORT) (imageBase + pImageExportDirectory->AddressOfNameOrdinals);

    ULONG i;
    for (i = 0; i < pImageExportDirectory->NumberOfNames; i++)
    {
        if (_stricmp(RoutineName, (PCHAR) imageBase + NamesArray[i]) == 0)
        {
            break;
        }
    }

    if (i > pImageExportDirectory->NumberOfNames - 1)
    {
        KiDebug("GetExportedRoutineByName Failed");
        *pFunctionAddress = 0;
        return FALSE;
    }

    ULONG Ordinal =  pImageExportDirectory->Base + NameOrdinalsArray[i];
    if (Ordinal > pImageExportDirectory->NumberOfFunctions -1 + pImageExportDirectory->Base)
    {
        KiDebug("GetExportedRoutineByName Failed");
        *pFunctionAddress = 0;
        return FALSE;
    }

    ULONG functionAddress = imageBase + FunctionsArray[NameOrdinalsArray[i]];

    KiDebug("Ordinal:%2d Hint:%2d Function:%s Entry Point:0x%08x", Ordinal, i, RoutineName, FunctionsArray[NameOrdinalsArray[i]]);
    if (!MmIsAddressValid((PVOID)functionAddress))
    {
        KiDebug("Address:0x%x is not Valid", functionAddress);
    }
    *pFunctionAddress = functionAddress;
    return TRUE ;
}

ULONG KiPE::GetNewImageBase()
{
    return mNewImageBase;
}
ULONG KiPE::GetOldImageBase()
{
    return mOldImageBase;
}
